package com.caissa.framework.starter.autoconfigure.reference;

import com.dianping.cat.Cat;
import com.dianping.cat.message.Message;
import com.dianping.cat.message.Transaction;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.aop.framework.ProxyFactoryBean;
import org.springframework.kafka.listener.AbstractMessageListenerContainer;
import org.springframework.kafka.listener.MessageListener;
import org.springframework.kafka.listener.MessageListenerContainer;

/**
 *
 * kafka 消费者埋点
 *
 * @auther wangwei
 * @date 2020/11/4
 */
@Aspect
public class KafkaConsumerCatAspect {

    private static final Log log = LogFactory.getLog(KafkaConsumerCatAspect.class);

    @Pointcut("execution(public * org.springframework.kafka.config.KafkaListenerContainerFactory.createListenerContainer(..))")
    private void anyCreateListenerContainer() {
    }

    @Around("anyCreateListenerContainer()")
    public Object wrapListenerContainerCreation(ProceedingJoinPoint pjp)
            throws Throwable {
        MessageListenerContainer listener = (MessageListenerContainer) pjp.proceed();
        if (listener instanceof AbstractMessageListenerContainer) {
            AbstractMessageListenerContainer container = (AbstractMessageListenerContainer) listener;
            Object someMessageListener = container.getContainerProperties()
                    .getMessageListener();
            if (someMessageListener == null) {
                if (log.isDebugEnabled()) {
                    log.debug("No message listener to wrap. Proceeding");
                }
            }
            else if (someMessageListener instanceof MessageListener) {
                container.setupMessageListener(createProxy(someMessageListener));
            }
            else {
                if (log.isDebugEnabled()) {
                    log.debug("ATM we don't support Batch message listeners");
                }
            }
        }
        else {
            if (log.isDebugEnabled()) {
                log.debug("Can't wrap this listener. Proceeding");
            }
        }
        return listener;
    }

    @SuppressWarnings("unchecked")
    Object createProxy(Object bean) {
        ProxyFactoryBean factory = new ProxyFactoryBean();
        factory.setProxyTargetClass(true);
        factory.addAdvice(
                new MessageListenerMethodInterceptor());
        factory.setTarget(bean);
        return factory.getObject();
    }

    static class MessageListenerMethodInterceptor<T extends MessageListener>
            implements MethodInterceptor {

        private static final Log log = LogFactory
                .getLog(MessageListenerMethodInterceptor.class);

        @Override
        public Object invoke(MethodInvocation invocation) throws Throwable {
            if (!"onMessage".equals(invocation.getMethod().getName())) {
                return invocation.proceed();
            }
            Object[] arguments = invocation.getArguments();
            Object record = record(arguments);
            if (record == null) {
                return invocation.proceed();
            }
            if (log.isDebugEnabled()) {
                log.debug("Wrapping onMessage call");
            }
            ConsumerRecord consumerRecord = (ConsumerRecord<?, ?>) record;
            String topic = record instanceof ConsumerRecord ? consumerRecord.topic() : "UNKNOW_TOPIC";
            Transaction transaction = Cat.newTransaction("kafka", new StringBuffer("consumer::").append(topic).toString());
            try{
                transaction.setStatus(Message.SUCCESS);
                return invocation.proceed();
            }
            catch (RuntimeException | Error e) {
                transaction.setStatus(e);
                throw e;
            }
            finally {
                transaction.complete();
            }
        }

        private Object record(Object[] arguments) {
            for (Object object : arguments) {
                if (object instanceof ConsumerRecord) {
                    return object;
                }
            }
            return null;
        }
    }
}


